{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e1feef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, Any, Tuple, Optional, Set\n",
    "import itertools\n",
    "from deepsignal3.utils.process_utils import complement_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "70a36f95",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_fasta(fasta_file: str) -> Dict[str, str]:\n",
    "    \"\"\"\n",
    "    读取 FASTA 文件，返回 {chrom: sequence} 字典，序列全转为大写。\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    fasta_file : str\n",
    "        参考基因组 FASTA 文件路径（支持 .fa, .fasta, .fa.gz）\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    ref_seq : dict\n",
    "        {chromosome_name: \"ATCG...\"}\n",
    "    \"\"\"\n",
    "    import gzip\n",
    "\n",
    "    ref_seq = {}\n",
    "    current_chrom = None\n",
    "    seq_lines = []\n",
    "\n",
    "    open_func = gzip.open if fasta_file.endswith('.gz') else open\n",
    "\n",
    "    with open_func(fasta_file, 'rt') as f:\n",
    "        for line in f:\n",
    "            line = line.strip()\n",
    "            if not line:\n",
    "                continue\n",
    "            if line.startswith('>'):\n",
    "                # 保存上一个染色体\n",
    "                if current_chrom is not None:\n",
    "                    ref_seq[current_chrom] = ''.join(seq_lines).upper()\n",
    "                # 开始新的染色体\n",
    "                current_chrom = line[1:].split()[0]  # 取 > 后的第一个词\n",
    "                seq_lines = []\n",
    "            else:\n",
    "                seq_lines.append(line)\n",
    "\n",
    "    # 保存最后一个\n",
    "    if current_chrom is not None:\n",
    "        ref_seq[current_chrom] = ''.join(seq_lines).upper()\n",
    "\n",
    "    return ref_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "306707d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref=load_fasta('/homeb/hpc/refs/genome/Ricinus.communis/GCF_019578655.1_ASM1957865v1_genomic.fna.fa')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "22a3be5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rev_seq_dict: Dict[str, str] = {}\n",
    "for chrom, seq in ref.items():\n",
    "    rev_seq_dict[chrom] = complement_seq(seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "8c38f36b",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines=30\n",
    "i=0\n",
    "bed_records = []\n",
    "with open('/homeb/hpc/refs/wgbs/ricinus_communis/ricinus_communis_1_val_1_bismark_bt2_pe.deduplicated.bismark.all_C.gz.bismark.zero.C.CHG.bed', \"r\") as rf:\n",
    "        for line in rf:\n",
    "            words = line.strip().split()\n",
    "            if len(words) < 6:\n",
    "                continue\n",
    "            chrom, pos_str, strand = words[0], words[1], words[5]\n",
    "            bed_records.append((chrom, int(pos_str), strand))\n",
    "            i+=1\n",
    "            if i>lines:\n",
    "                 break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c5690c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_motif(chrom, start, strand):\n",
    "    \"\"\"Extract 3bp motif centered at cytosine position from fasta\"\"\"\n",
    "    # 对于 BED 是 0-based，C 的位置 = start\n",
    "    ref_seq = rev_seq_dict[chrom] if strand == '-' else ref[chrom]\n",
    "    seq_str = str(ref_seq)\n",
    "    rev_start = len(seq_str) - 1 - start\n",
    "    # 检查边界\n",
    "    if start < 0 or start + 3 > len(seq_str):\n",
    "        return None\n",
    "    if strand=='+':\n",
    "        motif = seq_str[start : start + 3].upper()\n",
    "    else:\n",
    "        rev_start = len(seq_str) - 1 - start\n",
    "        motif = ref_seq[rev_start : rev_start + 3].upper()\n",
    "    return motif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "3c7ca53f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NC_063256.1 289 + CCG\n",
      "NC_063256.1 337 + CTG\n",
      "NC_063256.1 339 - CAG\n",
      "NC_063256.1 538 + CAG\n",
      "NC_063256.1 540 - CTG\n",
      "NC_063256.1 628 + CAG\n",
      "NC_063256.1 630 - CTG\n",
      "NC_063256.1 643 + CAG\n",
      "NC_063256.1 645 - CTG\n",
      "NC_063256.1 657 + CAG\n",
      "NC_063256.1 659 - CTG\n",
      "NC_063256.1 666 + CAG\n",
      "NC_063256.1 668 - CTG\n",
      "NC_063256.1 781 + CTG\n",
      "NC_063256.1 783 - CAG\n",
      "NC_063256.1 799 + CAG\n",
      "NC_063256.1 801 - CTG\n",
      "NC_063256.1 888 + CTG\n",
      "NC_063256.1 890 - CAG\n",
      "NC_063256.1 897 - CCG\n",
      "NC_063256.1 935 + CAG\n",
      "NC_063256.1 937 - CTG\n",
      "NC_063256.1 965 + CAG\n",
      "NC_063256.1 967 - CTG\n",
      "NC_063256.1 993 - CCG\n",
      "NC_063256.1 1085 + CTG\n",
      "NC_063256.1 1087 - CAG\n",
      "NC_063256.1 1096 - CCG\n",
      "NC_063256.1 1121 - CCG\n",
      "NC_063256.1 1133 + CTG\n",
      "NC_063256.1 1135 - CAG\n"
     ]
    }
   ],
   "source": [
    "for chrom, start, strand in bed_records:\n",
    "    motif = get_motif(chrom, start, strand)\n",
    "    print(chrom, start, strand, motif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "e92db640",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines=30\n",
    "i=0\n",
    "bed_records = []\n",
    "with open('/homeb/hpc/refs/wgbs/ricinus_communis/ricinus_communis_1_val_1_bismark_bt2_pe.deduplicated.bismark.all_C.gz.bismark.zero.C.CHH.bed', \"r\") as rf:\n",
    "        for line in rf:\n",
    "            words = line.strip().split()\n",
    "            if len(words) < 6:\n",
    "                continue\n",
    "            chrom, pos_str, strand = words[0], words[1], words[5]\n",
    "            bed_records.append((chrom, int(pos_str), strand))\n",
    "            i+=1\n",
    "            if i>lines:\n",
    "                 break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "42b7eb1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NC_063256.1 17 - CAA\n",
      "NC_063256.1 19 + CTT\n",
      "NC_063256.1 33 - CAA\n",
      "NC_063256.1 39 - CAA\n",
      "NC_063256.1 40 + CTT\n",
      "NC_063256.1 44 + CAA\n",
      "NC_063256.1 48 - CTT\n",
      "NC_063256.1 59 + CAT\n",
      "NC_063256.1 70 + CTA\n",
      "NC_063256.1 79 - CAA\n",
      "NC_063256.1 89 + CTT\n",
      "NC_063256.1 94 + CAT\n",
      "NC_063256.1 100 + CTT\n",
      "NC_063256.1 103 + CAA\n",
      "NC_063256.1 107 + CTC\n",
      "NC_063256.1 109 + CTC\n",
      "NC_063256.1 111 + CTT\n",
      "NC_063256.1 115 - CTA\n",
      "NC_063256.1 118 + CTT\n",
      "NC_063256.1 124 - CAA\n",
      "NC_063256.1 125 + CAT\n",
      "NC_063256.1 129 + CCA\n",
      "NC_063256.1 130 + CAA\n",
      "NC_063256.1 139 + CAA\n",
      "NC_063256.1 145 + CAT\n",
      "NC_063256.1 149 - CAA\n",
      "NC_063256.1 167 + CTT\n",
      "NC_063256.1 170 - CAA\n",
      "NC_063256.1 183 + CAT\n",
      "NC_063256.1 200 + CTC\n",
      "NC_063256.1 202 + CTA\n"
     ]
    }
   ],
   "source": [
    "for chrom, start, strand in bed_records:\n",
    "    motif = get_motif(chrom, start, strand)\n",
    "    print(chrom, start, strand, motif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0df366fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines=30\n",
    "i=0\n",
    "bed_records = []\n",
    "with open('/homeb/hpc/refs/wgbs/ricinus_communis/ricinus_communis_1_val_1_bismark_bt2_pe.deduplicated.bismark.all_C.gz.bismark.zero.C.CG.bed', \"r\") as rf:\n",
    "        for line in rf:\n",
    "            words = line.strip().split()\n",
    "            if len(words) < 6:\n",
    "                continue\n",
    "            chrom, pos_str, strand = words[0], words[1], words[5]\n",
    "            bed_records.append((chrom, int(pos_str), strand))\n",
    "            i+=1\n",
    "            if i>lines:\n",
    "                 break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "c9b8da86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NC_063256.1 54 + CGT\n",
      "NC_063256.1 55 - CGA\n",
      "NC_063256.1 284 + CGT\n",
      "NC_063256.1 285 - CGC\n",
      "NC_063256.1 290 + CGT\n",
      "NC_063256.1 291 - CGG\n",
      "NC_063256.1 307 + CGC\n",
      "NC_063256.1 308 - CGA\n",
      "NC_063256.1 719 + CGC\n",
      "NC_063256.1 720 - CGA\n",
      "NC_063256.1 895 + CGG\n",
      "NC_063256.1 896 - CGT\n",
      "NC_063256.1 982 + CGA\n",
      "NC_063256.1 983 - CGA\n",
      "NC_063256.1 991 + CGG\n",
      "NC_063256.1 992 - CGT\n",
      "NC_063256.1 998 + CGT\n",
      "NC_063256.1 999 - CGT\n",
      "NC_063256.1 1011 + CGA\n",
      "NC_063256.1 1012 - CGT\n",
      "NC_063256.1 1024 + CGA\n",
      "NC_063256.1 1025 - CGC\n",
      "NC_063256.1 1094 + CGG\n",
      "NC_063256.1 1095 - CGT\n",
      "NC_063256.1 1119 + CGG\n",
      "NC_063256.1 1120 - CGC\n",
      "NC_063256.1 1140 + CGC\n",
      "NC_063256.1 1141 - CGC\n",
      "NC_063256.1 1147 + CGG\n",
      "NC_063256.1 1148 - CGT\n",
      "NC_063256.1 1161 + CGT\n"
     ]
    }
   ],
   "source": [
    "for chrom, start, strand in bed_records:\n",
    "    motif = get_motif(chrom, start, strand)\n",
    "    print(chrom, start, strand, motif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deepsignal3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
